# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: zijin

import os
import datetime
import tempfile
import traceback

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
from itertools import combinations
import pickle as pkl
import glob

from coin.base.query_util import query_pta
from cc.appcoin2.strategy.basis_strat.support.daily_search.slack_helper import SlackHelper
from python.coin.strategy.basis_strat.analysis.report_og_log_analysis import extract_trading_info_from_og_log, get_start_date, get_end_date, get_all_strats_info

from absl import app, flags
FLAGS = flags.FLAGS


def report_strat(log_dict, strat, symbols):
  strat_info = []
  with ProcessPoolExecutor(max_workers=FLAGS.max_workers) as executor:
    futures = {}
    for key in log_dict:
      onlog = log_dict[key][0]
      if len(log_dict[key]) > 1:
        other_logs = log_dict[key][1:]
      else:
        other_logs = None
      fut = executor.submit(onlog.output_pnls_df, other_logs, symbols)
      futures[fut] = strat
    for f in as_completed(futures.keys()):
      try:
        res = f.result()
        if res:
          pnl_df = res[0]
          strat_info.append(pnl_df)
          print(f"{futures[f]['name']} report done")
        else:
          print(f"{futures[f]['name']} report empty")
      except Exception:
        import traceback
        traceback.print_exc()
  return strat_info


def get_strat_groups():
  # get model port and mm port strategies
  end_date = datetime.datetime.today() - datetime.timedelta(days=1)
  start_date = end_date - datetime.timedelta(days=1)
  pta_data = query_pta(
      start_date,
      end_date,
      business_units=['Coin'],
      pta_type='STRATEGY')
  mmports = sorted(set(pta_data.loc[pta_data['strategy_group'] == 'mm_portfolio']['strategy_name']))
  modelports = sorted(
      set(pta_data.loc[pta_data['strategy_group'] == 'model_portfolio']['strategy_name']))
  return mmports, modelports


def process_strat(strat, start, end, cpu):
  start_dt = datetime.datetime.strptime(start, '%Y%m%d')
  end_dt = datetime.datetime.strptime(end, '%Y%m%d')
  cur = start_dt
  first = True
  with ProcessPoolExecutor(max_workers=cpu) as executor:
    futures = {}
    while cur < end_dt:
      fut = executor.submit(extract_trading_info_from_og_log, strat['name'], strat['host'], cur.strftime('%Y%m%d'),
                            (cur + datetime.timedelta(days=1)).strftime('%Y%m%d'), strat['mea'], strat['symbols'], is_start=first)
      first = False
      futures[fut] = (strat, cur.strftime('%Y%m%d'))
      cur += datetime.timedelta(days=1)
    logs = []
    for f in as_completed(futures.keys()):
      try:
        logs.append((f.result(), futures[f][1]))
        print(f"{futures[f][0]['name']} {futures[f][0]['mea']} {futures[f][1]} extract done")
      except Exception:
        import traceback
        traceback.print_exc()
    log_dict = {}
    symbols = []
    for log, date in logs:
      strat_name = log.strategy_name
      mea = log.mea
      symbols.extend(log._sublogs.keys())
      if not (strat_name, mea) in log_dict:
        log_dict[(strat_name, mea)] = []
      if date == start:
        log_dict[(strat_name, mea)].insert(0, log)
      else:
        log_dict[(strat_name, mea)].append(log)
    symbols = set(symbols)
    strat_info = report_strat(log_dict, strat, symbols)
  return strat_info


def retrieve_strat_info(strategy_name_contains, strategy_type, trade_bases, trade_exchanges):
  start = get_start_date()
  end = get_end_date(start)
  infos = get_all_strats_info(
        strategy_name_contains=strategy_name_contains,
        strategy_type=strategy_type,
        trade_bases=trade_bases.split(',') if trade_bases else None,
        trade_exchanges=trade_exchanges.split(',')
        if trade_exchanges else None)

  strat_info = []
  strat_names = []
  days = (datetime.datetime.strptime(end, '%Y%m%d') -
          datetime.datetime.strptime(start, '%Y%m%d')).days
  cpu1 = max(FLAGS.max_workers // days, 1)
  cpu2 = max(FLAGS.max_workers // cpu1, 1)
  with ProcessPoolExecutor(max_workers=cpu1) as executor:
    futures = {}
    for strat in infos:
      fut = executor.submit(process_strat, strat, start, end, cpu2)
      futures[fut] = strat
    for f in as_completed(futures.keys()):
      try:
        pnls = f.result()
        if len(pnls) > 0:
          strat_info.extend(pnls)
          strat_names.append(futures[f]['name'])
          print(f"{futures[f]['name']} done")
        else:
          print(f"{futures[f]['name']} no pnl and position.")
      except Exception:
        import traceback
        traceback.print_exc()
  if len(strat_info) == 0:
    raise ValueError("Invalid strategy name./ No pnl for strategy.")
  return strat_info, strat_names


def get_end_of_interval_pos(p1, p2, col, interval):
  p1['ts'] = pd.to_datetime(p1['ts'])
  p2['ts'] = pd.to_datetime(p2['ts'])

  p1 = p1.resample(f'{interval}T', on='ts').last().reset_index(drop=True)
  p2 = p2.resample(f'{interval}T', on='ts').last().reset_index(drop=True)
  return p1[col], p2[col]


def intersection(lst1, lst2):
  lst3 = [value for value in lst1 if value in lst2]
  return lst3


def calculate_corrs(dfs1, dfs2, symbols, interval=60):
  res = []
  symbols = intersection(list(dfs1.keys()), list(dfs2.keys()))
  for symbol in tqdm(symbols):
    col = 'dollar_position'
    p1 = dfs1[symbol][['ts', col]].dropna()
    p2 = dfs2[symbol][['ts', col]].dropna()
    if p1.empty or p2.empty:
      continue
    df1, df2 = get_end_of_interval_pos(p1, p2, col, interval)
    df = pd.DataFrame({f'{col}1': df1, f'{col}2': df2})
    corr = df[f'{col}1'].corr(df[f'{col}2'])
    if not np.isnan(corr):
      res.append([symbol, corr])
  return res


def calculate_weighted_corr(pos_corrs, weights1, weights2):
  agg1 = 0
  agg2 = 0
  ws1 = sum(weights1.values())
  ws2 = sum(weights2.values())
  for pair in pos_corrs:
    symbol, corr = tuple(pair)
    agg1 += weights1[symbol] * corr / ws1
    agg2 += weights2[symbol] * corr / ws2
  return (agg1 + agg2) / 2


def calculate_twp(strat_names, strat_info):
  wdict = {}
  for name, pnls in zip(strat_names, strat_info):
    weights = {}
    for symbol in pnls.keys():
      df = pnls[symbol].dropna(subset=['dollar_position'])
      if df.empty:
        continue
      df['ts'] = pd.to_datetime(df['ts'])
      df = df.resample(f'{FLAGS.interval}T', on='ts').last().reset_index(drop=True)
      df['aot'] = df.index  # total amount of time
      twp = (df['aot'] * df['dollar_position'].abs()).sum() / df['aot'].sum()
      weights[symbol] = twp
    wdict[name] = weights
  return wdict


def common_start(values):
  if len(values) != 0:
    start = values[0]
    while start:
      if all(value.startswith(start) for value in values):
        return start
      start = start[:-1]
  return ""


def plot_colormap(matrix, strat_names, tmp_output_dir):
  n = len(strat_names)
  cs = common_start(strat_names)
  simple_names = [name.replace(cs, '*') for name in strat_names]
  size = max(8, min(16, n))
  fig, ax = plt.subplots(figsize=(size, size))

  matrix = np.array(matrix)
  ax.matshow(matrix, cmap=plt.cm.Blues)

  ax.set_xticks(list(range(n)))
  ax.set_yticks(list(range(n)))
  ax.set_xticklabels(simple_names)
  ax.set_yticklabels(simple_names)
  plt.setp(ax.get_xticklabels(), rotation=90, horizontalalignment='right')

  for i in range(n):
    for j in range(n):
      c = matrix[j][i]
      ax.text(i, j, str(round(c, 4)), va='center', ha='center')

  figname = os.path.join(tmp_output_dir, f'test_{FLAGS.interval}.png')
  plt.title(f'{cs}(*)')
  plt.tight_layout()
  fig.savefig(figname)
  return figname


def get_strategy_correlation(strat_info, strat_names, tmp_output_dir):
  n = len(strat_names)
  matrix = np.ones((n, n))
  wdict = calculate_twp(strat_names, strat_info)
  for idx1, idx2 in combinations(list(range(len(strat_names))), 2):
    n1 = strat_names[idx1]
    n2 = strat_names[idx2]
    pos_corrs = calculate_corrs(strat_info[idx1], strat_info[idx2], FLAGS.interval)
    corr = calculate_weighted_corr(pos_corrs, wdict[n1], wdict[n2])
    matrix[idx1][idx2] = matrix[idx2][idx1] = corr
  figname = plot_colormap(matrix, strat_names, tmp_output_dir)
  if FLAGS.report_slack:
    try:
      sh = SlackHelper('.', FLAGS.slack_channel)
      sh.send_slack_noti(f'{figname} done', figname, True)
    except BaseException:
      traceback.print_exc()


def retrieve_sim_data():
  strat_names = sorted(os.listdir(FLAGS.sim_strat_path))
  strat_paths = [
      glob.glob(
          os.path.join(
              FLAGS.sim_strat_path,
              name,
              'stats',
              '*.gzip'))[0] for name in strat_names]
  strat_info = []
  for strat in strat_paths:
    strat = pd.read_csv(strat, compression='gzip')
    dfs = {}
    for group in strat.groupby(['symbol']):
      df = group[1][['timestamp.1', 'open_cash']].rename(
          columns={'timestamp.1': 'ts', 'open_cash': 'dollar_position'})
      df['multiplier'] = 1
      dfs[group[0]] = df
    strat_info.append(dfs)
  return strat_info, strat_names


def main(argv):
  tmp_output_dir = tempfile.mkdtemp()
  print(f'Output saved in {tmp_output_dir}')
  if FLAGS.selected_live_strat_names:
    strat_info = []
    strat_names = sorted(FLAGS.selected_live_strat_names.split(','))
    for strat in strat_names:
      single_info, _ = retrieve_strat_info(strat, '', '', '')
      strat_info.extend(single_info)
  elif FLAGS.sim_strat_path:
    strat_info, strat_names = retrieve_sim_data()
  else:
    strat_info, strat_names = retrieve_strat_info(FLAGS.strategy_contains, '', '', '')
    tmp_info = []
    tmp_names = []
    for info, name in sorted(zip(strat_info, strat_names), key=lambda x: x[1]):
      tmp_info.append(info)
      tmp_names.append(name)
    strat_info = tmp_info
    strat_names = tmp_names
  pkl.dump((strat_info, strat_names), open(os.path.join(tmp_output_dir, 'data.pkl'), 'wb'))

  #strat_info, strat_names = pkl.load(open(os.path.join('/tmp/tmpf4lbikui', 'data.pkl'), 'rb'))
  assert(len(strat_info) > 0)
  get_strategy_correlation(strat_info, strat_names, tmp_output_dir)


if __name__ == '__main__':
  flags.DEFINE_string('start_date', '20220320', '')
  flags.DEFINE_string('end_date', '20220330', '')
  flags.DEFINE_string('strategy_contains', 'mm_portfolio', '')
  flags.DEFINE_string(
      'selected_live_strat_names',
      'model_portfolio_binance_perpetual_v1,model_portfolio_binance_perpetual_v1_2',
      '')
  flags.DEFINE_string(
      'sim_strat_path',
      '/remote/iosg/home-3/yuxuan/coin/sim_live/model_port_prod3',
      '')
  flags.DEFINE_integer('interval', 10, 'get end of X minute position')
  flags.DEFINE_string('from_mail', 'basis_strategy@joomo.io', '')
  flags.DEFINE_string('to_mail', 'zijin@prestolabs.io', '')
  flags.DEFINE_string('error_mail', '', '')
  flags.DEFINE_bool('report_slack', True, '')
  flags.DEFINE_string('slack_channel', 'zijin_test', '')
  flags.DEFINE_integer('max_workers', 8, '')
  flags.DEFINE_float('shift_start_date_hours', 0, '')
  flags.DEFINE_float('shift_funding_fee_start_date_hours', 1, '')
  flags.DEFINE_bool('dump_to_file', False, '')
  flags.DEFINE_bool('light_ver', False, 'light version only has pnl plot')
  flags.DEFINE_string('bbo_cache_root',
                      '/remote/iosg/data-2/buckets/feed.derived.interval_h5/coin/main/PT1M/', '')
  flags.DEFINE_bool('include_funding_rate', True, '')
  app.run(main)
